
#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/slab.h>
#include <linux/kthread.h>
#include <linux/delay.h>
#include <linux/time.h>
#include "hisi_l0_mem_pool.h"

#define THREAD_NUM 8
#define ALLOC_TIMES 100
#define TEST_BLOCK_SIZE 4096

static struct task_struct *threads[THREAD_NUM];
static u64 thread_times[THREAD_NUM];

// 线程工作函数
static int kmalloc_worker(void *data) {
    int thread_id = *(int *)data;
    ktime_t start, end;
    s64 delta;
    int i;

    start = ktime_get();
    for (i = 0; i < ALLOC_TIMES; i++) {
        void *ptr = kmalloc(TEST_BLOCK_SIZE, GFP_KERNEL);
        if (!ptr) {
            printk(KERN_ERR "kmalloc failed in thread %d\n", thread_id);
            return -ENOMEM;
        }
        kfree(ptr);
    }
    end = ktime_get();

    delta = ktime_to_ns(ktime_sub(end, start));
    thread_times[thread_id] = delta;

    printk(KERN_INFO "Thread %d: %llu ns\n", thread_id, delta);
    return 0;
}

// 线程工作函数
static int l0_kmalloc_worker(void *data) {
    int thread_id = *(int *)data;
    ktime_t start, end;
    s64 delta;
    int i;
    //unsigned long alloc_size = 0;
	
    start = ktime_get();
    for (i = 0; i < ALLOC_TIMES; i++) {
	//printk(KERN_ERR "l0_kmalloc_worker alloc size %lu\n", alloc_size);
	//void *pool_handle1 = NULL;
        void* ptr = l0_kmalloc(TEST_BLOCK_SIZE, 0);
        if (!ptr) {
            printk(KERN_ERR "kmalloc failed in thread %d\n", thread_id);
            return -ENOMEM;
        }
	//alloc_size += TEST_BLOCK_SIZE;
        l0_kfree(ptr);
        //alloc_size -= TEST_BLOCK_SIZE;
    }
    end = ktime_get();

    delta = ktime_to_ns(ktime_sub(end, start));
    thread_times[thread_id] = delta;

    printk(KERN_INFO "Thread %d: %llu ns\n", thread_id, delta);
    return 0;
}


// 模块初始化
static int __init kmalloc_test_init(void) {
    int i;
    int thread_ids[THREAD_NUM];

    printk(KERN_INFO "Starting kmalloc performance test\n");

    for (i = 0; i < THREAD_NUM; i++) {
        thread_ids[i] = i;
        threads[i] = kthread_run(l0_kmalloc_worker, &thread_ids[i], "kmalloc_test/%d", i);
        if (IS_ERR(threads[i])) {
            printk(KERN_ERR "Failed to create thread %d\n", i);
            return PTR_ERR(threads[i]);
        }
    }

    return 0;
}

static void safe_stop_thread(struct task_struct *thread) {
    if (!IS_ERR_OR_NULL(thread)) {
        kthread_stop(thread);
        thread = NULL;
    }
}

// 模块退出
static void __exit kmalloc_test_exit(void) {
    //int i;
    u64 total = 0;

    // for (i = 0; i < THREAD_NUM; i++) {
    //     if (threads[i])
    //         safe_stop_thread(threads[i]);
    //     total += thread_times[i];
    // }

    printk(KERN_INFO "Average time per kmalloc: %llu ns\n",
           total / (THREAD_NUM * ALLOC_TIMES));
    printk(KERN_INFO "kmalloc test module unloaded\n");
}

module_init(kmalloc_test_init);
module_exit(kmalloc_test_exit);
MODULE_LICENSE("GPL");
MODULE_AUTHOR("Your Name");
MODULE_DESCRIPTION("Kernel module to test kmalloc performance");
